// Yue Marvin Tao
// 3/24/2022

package mult 

import chisel3._
import chisel3.util._

class SAMInputBundle(val bits: Int) extends Bundle {
  val load = Input(Bool())
  val A = Input(UInt(bits.W))
  val B = Input(UInt(bits.W))
}

class SAMOutputBundle(val bits: Int) extends Bundle {
  val P = Output(UInt((2 * bits).W))
  val M = Output(UInt((2 * bits).W))
  val c = Output(UInt(bits.W))
}

class ShiftAddMul(val bits:Int) extends Module {
  val in  = IO(new SAMInputBundle(bits))
  val out = IO(new SAMOutputBundle(bits))
  
  val move = RegNext(VecInit(Seq.fill(2 * bits)(0.U(1.W))))
  val prod = RegNext(VecInit(Seq.fill(2 * bits)(0.U(1.W))))

  val carry = Wire(Vec(2 * bits + 1, UInt(1.W))) 
  val count = RegNext(0.U(bits.W))

  for (i <- 0 until 2 * bits + 1) {
    carry(i) := 0.U
  }

  when(in.load) {
    for (i <- 0 until bits) {
      move(i) := in.B(i)
    }
  }.elsewhen(count < (bits * 2).U){
    // A[d] && prod += move
    when(in.A(count).asBool) {
      for (i <- 0 until 2 * bits) {
        val p = prod(i)
        val m = move(i)
        val c = carry(i)
        val cn = carry(i+1)
        p := p ^ m ^ c 
        cn := (p & m) | (m & c) | (c & p) 
      }
    }.otherwise {
      prod := prod
    }

    // move << 1
    for (i <- 0 until (2 * bits - 1)) {
      move(i+1) := move(i)
    }
  
    count := count + 1.U
  }

  out.M := move.asUInt
  out.P := prod.asUInt
  out.c := count
}


object Top extends App {
  (new chisel3.stage.ChiselStage).emitVerilog(new ShiftAddMul(4))
}